package com.byq.mcpanus.service;

import com.byq.mcpanus.model.Conversation;
import com.byq.mcpanus.model.McpServer;
import com.byq.mcpanus.model.DomainServer;
import com.byq.mcpanus.model.McpDomain;
import com.byq.mcpanus.repository.McpServerRepository;
import com.byq.mcpanus.repository.DomainServerRepository;
import com.byq.mcpanus.repository.McpDomainRepository;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import jakarta.annotation.PostConstruct;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.DefaultChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.model.Media;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.MimeType;


import java.lang.reflect.Field;
import java.util.*;

@Service
public class ChatService {
    @Autowired 
    private ChatClient.Builder chatClientBuilder;
    @Autowired 
    private McpServerRepository mcpServerRepository;
    @Autowired 
    private DomainServerRepository domainServerRepository;
    @Autowired 
    private McpDomainRepository mcpDomainRepository;
    
    private ObjectMapper objectMapper = new ObjectMapper();
    private Map<String, ChatClient> domainChatClients = new HashMap<>();
    Map<String, Conversation> conversations = new HashMap<>();

    @Autowired ConversationService cservice;

    String systemMessage= """
            You are a doctor to help people keeping healthy. 
            A question is given which is enclosed by ```.
            Please give your response in chinese
            """;
    public void initConversation(){
        List<Conversation> cs = cservice.getAll();
        for(Conversation c :  cs){
            conversations.put(c.getId(), c);
        }


    }
    @PostConstruct
    public void init() {
        List<McpDomain> domains = mcpDomainRepository.findAll();
        for (McpDomain domain : domains) {
            initDomainChatClient(domain);
        }
        initConversation();
    }

    private void initDomainChatClient(McpDomain domain) {
        try {
            List<McpSyncClient> mcpSyncClients = new ArrayList<>();
            // 获取该域下的所有服务器
            List<DomainServer> domainServers = domainServerRepository.findByDomainId(domain.getId());
            for (DomainServer domainServer : domainServers) {
                // McpServer.id为Long，DomainServer.serverId为String，需转换
                Optional<McpServer> mcpServerOpt = mcpServerRepository.findById(Long.valueOf(domainServer.getServerId()));
                if (mcpServerOpt.isPresent()) {
                    McpServer mcpServer = mcpServerOpt.get();
                    String config = mcpServer.getConfig();
                    // 解析 JSON 配置
                    JsonNode configNode = objectMapper.readTree(config);
                    JsonNode mcpServers = configNode.get("mcpServers");
                    if (mcpServers != null) {
                        for (JsonNode server : mcpServers) {
                            String transport = server.get("transport").asText();
                            McpSyncClient mcpClient;
                            if (transport.equals("stdio")) {
                                String command = server.get("command").asText();
                                List<String> args = new ArrayList<>();
                                if (server.has("args")) {
                                    for (JsonNode arg : server.get("args")) {
                                        args.add(arg.asText());
                                    }
                                }
                                ServerParameters stdioParams = ServerParameters.builder(command)
                                        .args(args.toArray(new String[0]))
                                        .build();
                                mcpClient = McpClient.sync(new StdioClientTransport(stdioParams)).build();
                            } else if (transport.equals("sse")) {
                                String url = server.get("url").asText();
                                mcpClient = McpClient.sync(new HttpClientSseClientTransport(url)).build();
                            } else {
                                throw new RuntimeException("unsupported mcp transport " + transport);
                            }
                            McpSchema.InitializeResult init = mcpClient.initialize();
                            mcpSyncClients.add(mcpClient);
                        }
                    }
                }
            }
            SyncMcpToolCallbackProvider provider = new SyncMcpToolCallbackProvider(mcpSyncClients);
            clear(chatClientBuilder);
            ChatClient chatClient = chatClientBuilder.defaultSystem("You are useful assistant and can operate your browser.")
                    .defaultTools(provider)
                    .defaultAdvisors(new MessageChatMemoryAdvisor(new InMemoryChatMemory()))
                    .build();
            domainChatClients.put(domain.getId(), chatClient);
        } catch (Exception e) {
            // 记录错误但不中断初始化
            System.err.println("Failed to initialize chat client for domain: " + domain.getId() + ", error: " + e.getMessage());
        }
    }
    void clear(Object chatClientBuilder){
        try {
            Field field = chatClientBuilder.getClass().getDeclaredField("defaultRequest");
            field.setAccessible(true); // 允许访问 protected/private 字段
            DefaultChatClient.DefaultChatClientRequestSpec req = (DefaultChatClient.DefaultChatClientRequestSpec) field.get(chatClientBuilder);
            req.getFunctionCallbacks().clear();
        } catch (Exception e) {
            throw new RuntimeException("无法获取 defaultChatRequest 字段", e);
        }
    }

    public String ask(String question, String cid,  String domain) throws Exception {
        Conversation con = null;

        if(cid != null) {
            con = conversations.get(cid);
            domain = con.getDomain();
        }else{
            con = new Conversation(UUID.randomUUID().toString());
            con.setDomain(domain);
            conversations.put(con.getId(), con);
            cid = con.getId();
            SystemMessage sm = new SystemMessage(systemMessage);
            con.getMessages().add(sm);
        }
        ChatClient chatClient = domainChatClients.get(domain);
        if (chatClient == null) {
            throw new Exception("No chat client found for domain: " + domain);
        }



        UserMessage um = new UserMessage(question);
        con.getMessages().add(um);

        List<Message> messages = con.getMessages();
        ChatClient.CallResponseSpec aws = chatClient.prompt().messages(messages)
                .call();

        String cont = aws.content();

        AssistantMessage am = new AssistantMessage(cont);


        con.getMessages().add(am);

        this.cservice.save(con);

        return cont;

    }

    // 新增支持图片base64的方法
    public String ask(String question, String cid, String domain, String imageBase64) throws Exception {
        Conversation con = null;
        if(cid != null) {
            con = conversations.get(cid);
            domain = con.getDomain();
        } else {
            con = new Conversation(UUID.randomUUID().toString());
            con.setDomain(domain);
            conversations.put(con.getId(), con);
            cid = con.getId();
            SystemMessage sm = new SystemMessage(systemMessage);
            con.getMessages().add(sm);
        }
        ChatClient chatClient = domainChatClients.get(domain);
        if (chatClient == null) {
            throw new Exception("No chat client found for domain: " + domain);
        }

        UserMessage um;
        if (imageBase64 != null && !imageBase64.isEmpty()) {
            String base64 = imageBase64;
            String mimeType = "image/jpeg";
            // 兼容 data:image/png;base64,... 或 data:image/webp;base64,...

            Media media = Media.builder()
                    .data(base64)
                    .mimeType(MimeType.valueOf(mimeType))
                    .build();
            um = new UserMessage(question, List.of(media));
        } else {
            um = new UserMessage(question);
        }
        con.getMessages().add(um);

        List<Message> messages = con.getMessages();
        ChatClient.CallResponseSpec aws = chatClient.prompt().messages(messages).call();

        String cont = aws.content();
        AssistantMessage am = new AssistantMessage(cont);
        con.getMessages().add(am);
        this.cservice.save(con);

        return cont;
    }
} 